# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import time
from pydantic import TypeAdapter
from tortoise.queryset import Q
from tortoise.functions import Count
from hypertext import PagingResult
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.enums.models import ModelsEnum
from common.enums.notice import NoticeEnum
from common.models.users import UserModel
from common.models.dev import DevModelsModel
from common.models.agents import AgentKnowModel
from common.postgres.public import PgKnowledgeModel
from plugins.msg.driver import MsgDriver
from apps.admin.schemas.agent.know import manage_schema as schema


class KnowManageService:

    @classmethod
    async def lists(cls, params: schema.KnowSearchIn) -> PagingResult[schema.KnowListVo]:
        """
        知识库列表。

        Args:
            params (schema.KnowSearchIn): 知识库查询参数。

        Returns:
            PagingResult[schema.KnowListVo]: 知识库分页列表Vo。

        Author:
            zero
        """
        # 查询条件
        where = AgentKnowModel.build_search({
            "%like%": ["name"],
            "=": ["is_disable"],
        }, params.__dict__)

        where.append(Q(is_delete=0))
        if params.user:
            condition = [Q(nickname__icontains=params.user) | Q(sn=params.user)]
            users = await UserModel.filter(*condition).filter(is_delete=0).limit(200).values("id")
            ids = [item["id"] for item in users] or [0]
            where.append(Q(user_id__in=ids))

        # 查询数据
        _model = AgentKnowModel.filter(*where).order_by("-id")
        _pager = await AgentKnowModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size,
            fields=AgentKnowModel.without_field("sort,is_delete,delete_time")
        )

        # 查询用户
        user_ids = [int(item["user_id"]) for item in _pager.lists if item["user_id"]]
        users: dict = await UserModel.queryUserDictsById(user_ids)

        # 查拆分模型
        splits_list = [item["splits_model"] for item in _pager.lists if item["splits_model"]]
        splits_model = await DevModelsModel.get_model_alias(splits_list, ModelsEnum.TYPE_CHAT)

        # 查向量模型
        vector_list = [item["vector_model"] for item in _pager.lists if item["vector_model"]]
        vector_model = await DevModelsModel.get_model_alias(vector_list, ModelsEnum.TYPE_VECTOR)

        # 查询数据量
        know_ids = [int(item["id"]) for item in _pager.lists if item["id"]]
        total_query = await (PgKnowledgeModel
                             .filter(is_delete=0)
                             .filter(know_id__in=know_ids)
                             .annotate(count=Count("know_id"))
                             .group_by("know_id")
                             .values_list("know_id", "count"))
        total_data = {k: v for k, v in total_query}

        # 处理数据项
        _lists = []
        for item in _pager.lists:
            user = users.get(item["user_id"], {})
            vModel = vector_model.get(item["vector_model"], {})
            sModel = splits_model.get(item["splits_model"], {})

            _vector_model = vModel.get("name") + ":" + vModel["model"] if vModel else ""
            _splits_model = sModel.get("name") + ":" + sModel["model"] if sModel else ""
            if vModel.get("name") == vModel["model"]:
                _vector_model = vModel.get("name")
            if sModel.get("name") == sModel["model"]:
                _splits_model = sModel.get("name")

            item["logo"] = await UrlUtil.to_absolute_url(item["logo"])
            item["vector_model"] = _vector_model
            item["splits_model"] = _splits_model
            item["train_total"] = total_data.get(item["id"], 0)
            item["user"] = schema.KnowListVo.User(
                sn=user.get("sn", "").strip(),
                mobile=user.get("mobile", ""),
                nickname=user.get("nickname", ""),
                avatar=await UrlUtil.to_absolute_url(user.get("avatar", ""))
            )

            _lists.append(TypeAdapter(schema.KnowListVo).validate_python(item))

        _pager.lists = _lists
        return _pager

    @classmethod
    async def detail(cls, id_: int) -> schema.KnowDetailVo:
        """
        知识库详情。

        Args:
            id_ (int): 知识库ID。

        Returns:
            schema.KnowDetailVo: 知识库详情Vo。

        Author:
            zero
        """
        # 查找知识库
        know = await AgentKnowModel.filter(id=id_, is_delete=0).get()

        # 返回信息
        return schema.KnowDetailVo(
            id=know.id,
            code=know.code,
            logo=await UrlUtil.to_absolute_url(know.logo),
            name=know.name,
            intro=know.intro,
            vector_model=know.vector_model,
            splits_model=know.splits_model,
            is_disable=know.is_disable,
            create_time=TimeUtil.timestamp_to_date(know.create_time),
            update_time=TimeUtil.timestamp_to_date(know.update_time)
        )

    @classmethod
    async def delete(cls, id_: int, reason: str) -> None:
        """
        知识库删除。

        Args:
            id_ (int): 知识库ID。
            reason (str): 原因信息。

        Author:
            zero
        """
        # 查找知识库
        know = await AgentKnowModel.filter(id=id_, is_delete=0).get()

        # 删除知识库
        await AgentKnowModel.filter(id=id_, is_delete=0).update(
            is_delete=1,
            delete_time=int(time.time())
        )

        # 发送用户通知
        await MsgDriver.send(scene=NoticeEnum.KNOW_DELETE, params={
            "user_id": int(know.user_id),
            "code": know.code,
            "name": know.name,
            "reason": reason
        })

    @classmethod
    async def disable(cls, id_: int, reason: str) -> None:
        """
        知识库禁用。

        Args:
            id_ (int): 知识库ID。
            reason (str): 原因信息。

        Author:
            zero
        """
        # 查找知识库
        know = await AgentKnowModel.filter(id=id_, is_delete=0).get()

        # 调整状态值
        status = 0 if know.is_disable else 1
        await AgentKnowModel.filter(id=id_, is_delete=0).update(
            is_disable=int(status),
            update_time=int(time.time())
        )

        # 发送用户通知
        scene = NoticeEnum.KNOW_RESTORE if status == 0 else NoticeEnum.KNOW_DISABLE
        await MsgDriver.send(scene=scene, params={
            "user_id": int(know.user_id),
            "code": know.code,
            "name": know.name,
            "reason": reason
        })
